# -*- coding: utf-8 -*-
"""
@Time    : 2024/8/23 18:50 
@Author  : ZhangShenao 
@File    : 9.不支持工具调用的大模型解决方案.py
@Desc    : 不支持函数调用的大模型解决方案

针对不支持工具调用的大模型,可以通过编写特定的Prompt
让大模型按照特定的结构化格式生成内容,并且按照对应的格式进行输出解析
这样可以模拟工具调用的能力

但是这种方式存在一定的不确定性,无法保证LLM每次的生成都满足特定的格式
而且Prompt也会比较复杂,会消耗更多的Token
"""
from typing import Optional

import dotenv
from langchain_community.chat_models import ChatZhipuAI
from langchain_community.tools import GoogleSerperRun
from langchain_community.utilities import GoogleSerperAPIWrapper
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough, RunnableConfig
from langchain_core.tools import render_text_description_and_args
from pydantic.v1 import BaseModel, Field

from gaode_weather_tool import GaodeWeatherTool


class SerperArgSchema(BaseModel):
    query: str = Field(description="执行搜索的查询语句")


# 定义工具调用请求的参数格式
class ToolCallRequest(BaseModel):
    """工具调用请求的参数格式"""

    name: str = Field(description="需要调用的工具名称")
    args: dict = Field(description="工具调用的参数字典，key表示参数名，value表示具体的参数值")


# 加载环境变量
dotenv.load_dotenv()

# 创建LLM
llm = ChatZhipuAI(model="glm-4-airx", temperature=0)

# 创建工具列表
serper_tool = GoogleSerperRun(
    name="google_serper",  # 指定工具名称
    # 指定工具描述
    description=(
        "一个低成本的谷歌搜索API。"
        "当你需要回答相关的新闻或时事问题时，可以调用该工具。"
        "该工具的输入是执行搜索的查询语句。"
    ),
    args_schema=SerperArgSchema,  # 指定输入参数格式
    api_wrapper=GoogleSerperAPIWrapper(),  # 指定Google API包装器
)
weather_tool = GaodeWeatherTool()
tools = {
    weather_tool.name: weather_tool,
    serper_tool.name: serper_tool
}

# 定义Json输出解析器格式
parser = JsonOutputParser(pydantic_object=ToolCallRequest)

# 构造Prompt
system_prompt = """你是一个功能强大的聊天机器人，可以调用工具来解答用户的问题。下面是可用的工具列表：

{tools_desc}

"""

human_prompt = """生成的内容需要严格遵守以下格式：
{format_instruction}

用户提问：{query}
"""

# 通过render_text_description_and_args工具函数,可以快速生成工具列表的描述和参数信息
prompt = ChatPromptTemplate.from_messages([
    ("system", system_prompt),
    ("human", human_prompt)
]).partial(
    tools_desc=render_text_description_and_args([tool for tool in tools.values()]),
    format_instruction=parser.get_format_instructions()
)


def call_tool(req: ToolCallRequest, config: Optional[RunnableConfig] = None):
    """实现工具调用"""

    tool = tools.get(req.get("name"))
    return tool.invoke(req.get("args"), config=config)


# 构造Chain
chain = ({"query": RunnablePassthrough()}
         | prompt
         | llm
         | parser
         | RunnablePassthrough.assign(output=call_tool)
         )

# 调用chain,获取结果
print(chain.invoke("北京今天天气怎么样？"))
